# Copyright (c) Microsoft Corporation.
# Licensed under the MIT License.

"""
Online service script for Raymedia (SZ300251) using Qlib's OnlineTool.
- First run trains a model and marks it as the online model.
- Daily run updates online predictions based on the latest data.

Usage examples:
  python predict_raymedia_price.py first_train
  python predict_raymedia_price.py update_online_pred
  python predict_raymedia_price.py main --experiment_name="raymedia_online"

Optional args:
  --provider_uri "~/.qlib/qlib_data/cn_data"
  --region cn
  --instrument csi300  # default to CSI300 universe for robustness
"""

import copy
import fire
import qlib
from qlib.constant import REG_CN
from qlib.model.trainer import task_train
from qlib.workflow.online.utils import OnlineToolR
from qlib.tests.config import CSI300_GBDT_TASK


def _set_instruments_safely(task_cfg: dict, instruments: str):
    """Try to set instruments in common locations of a task config.
    This keeps original structure when keys are missing.
    """
    try:
        # Typical path: dataset -> kwargs -> handler -> kwargs -> instruments
        task_cfg["dataset"]["kwargs"]["handler"]["kwargs"]["instruments"] = instruments
        return
    except Exception:
        pass
    try:
        # Alternate: dataset -> handler -> kwargs -> instruments
        task_cfg["dataset"]["handler"]["kwargs"]["instruments"] = instruments
        return
    except Exception:
        pass
    try:
        # Fallback flat key
        task_cfg["instruments"] = instruments
    except Exception:
        pass


class OnlineRaymediaService:
    def __init__(
        self,
        provider_uri: str = "~/.qlib/qlib_data/cn_data",
        region: str = REG_CN,
        experiment_name: str = "raymedia_online",
        instrument: str = "csi300",
    ):
        # Initialize Qlib with your data dir and region
        qlib.init(provider_uri=provider_uri, region=region)
        self.experiment_name = experiment_name
        self.online_tool = OnlineToolR(self.experiment_name)
        # Prepare a default task config and set instruments (default: csi300 for robust availability)
        self.task_config = copy.deepcopy(CSI300_GBDT_TASK)
        _set_instruments_safely(self.task_config, instrument)

    def first_train(self):
        """Train a model once and set it as the online model."""
        rec = task_train(self.task_config, experiment_name=self.experiment_name)
        self.online_tool.reset_online_tag(rec)  # set to online model
        print(f"Model trained and set online under experiment: {self.experiment_name}")

    def update_online_pred(self):
        """Update online predictions based on latest data (run daily)."""
        self.online_tool.update_online_pred()
        print("Online predictions updated.")

    def main(self):
        """Run full process: train-then-update (for quick bootstrap)."""
        self.first_train()
        self.update_online_pred()


if __name__ == "__main__":
    fire.Fire(OnlineRaymediaService)